import torch
from torch import nn, optim
from torch.utils.data import DataLoader, TensorDataset

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("device:", device)

X = torch.randn(1000, 10)
y = (X.sum(dim=1) > 0).long()
ds = TensorDataset(X, y)
loader = DataLoader(ds, batch_size=32, shuffle=True)

model = nn.Sequential(nn.Linear(10,32), nn.ReLU(), nn.Linear(32,2)).to(device)
loss_fn = nn.CrossEntropyLoss()
opt = optim.SGD(model.parameters(), lr=0.1)

for epoch in range(5):
    total = 0.0
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        pred = model(xb)
        loss = loss_fn(pred, yb)
        opt.zero_grad()
        loss.backward()
        opt.step()
        total += loss.item() * xb.size(0)
    print(f"epoch {epoch} loss {total/len(ds):.4f}")

torch.save(model.state_dict(), "model.pth")
